# 用于数据标注任务的Web服务器
from typing import List, Tuple, Dict
import os
from flask import Flask, request, jsonify, send_from_directory, render_template

print("当前工作目录：", os.getcwd())
class DhlpWeb(object):
    ds_fn = '../../work/datasets/v1/ds.txt'
    labeling_prog_fn = '../../work/datas/dhlp/labeling_prog.txt'
    base_folder = '../../work/datas/dhlp'
    did, rid = None, None
    gt_labels = {}
    labeling_progs = {} # 保存标注进度

    def __init__(self):
        self.name = 'apps.dhlp.dhlp_web.DhlpWeb'

    @staticmethod
    def get_direct_sub_folders(base_folder:str) -> List:
        entries = os.listdir(base_folder)
        return [entry for entry in entries if os.path.isdir(os.path.join(base_folder, entry))]

    @staticmethod
    def generate_labeling_prog(base_folder:str) -> None:
        '''
        生成指定目录下设备＋记录标注进度初始文件，内容为：
        设备编号+记录编号:0/1 # 其中0---未标注; 1---已标注;
        注：此方法只需在每次标注任务前运行一次
        '''
        devs = DhlpWeb.get_direct_sub_folders(base_folder)
        # devs = sorted(devs)
        with open(DhlpWeb.labeling_prog_fn, 'w', encoding='utf-8') as wfd:
            for dev in devs:
                recs = DhlpWeb.get_direct_sub_folders(f'{base_folder}/{dev}')
                # recs = sorted(recs)
                for rec in recs:
                    print(f'{dev}+{rec}:0')
                    wfd.write(f'{dev}+{rec}:0\n')

    @staticmethod
    def get_curr_unlabel(prog_fn:str, progs:Dict) -> Tuple[str, str]:
        '''
        从labeling_prog_fn文件中找到第一个未标注的记录，返回设备编号和记录编号
        '''
        if not progs: # 如果其中不包含任何元素则进行初始化
            with open(prog_fn, 'r', encoding='utf-8') as rfd:
                for row in rfd:
                    row = row.strip()
                    arrs0 = row.split(':')
                    arrs1 = arrs0[0].split('+')
                    state = arrs0[1]
                    progs[arrs0[0]] = state
                    # did, rid = arrs1[0], arrs1[1]
        keys = sorted(progs.keys())
        for key in keys:
            if progs[key] == '0':
                arrs0 = key.split('+')
                did, rid = arrs0[0], arrs0[1]
                break
        return did, rid

    @staticmethod
    def get_gt_labels(ds_fn:str) -> Dict:
        '''
        根据ds.txt文件中内容，生成如下字典：
        {
            'pt_fn': {
                'pt_fn': **,
                'did': ...,
                'rid': ...,
                'label': 0/1, # 0---正常记录; 1---雷电记录
                'pos': 0/1/2/3, # 0---正常帧; 1---雷电前10分钟; 2---雷电前10～20分钟； 3---雷电前20～30分钟；
                'cls': 1/2/3, # 1---可预测10分钟内发生雷电的帧；2---可预测10～20分钟发生雷电的帧；3---可预测20～30分钟发生雷电的帧；
            }
        }
        '''
        gt_labels = {}
        with open(ds_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                gt_labels[arrs0[0]] = {
                    'pt_fn': arrs0[0],
                    'did': arrs0[1],
                    'rid': arrs0[2],
                    'label': arrs0[3],
                    'pos': arrs0[4],
                    'cls': arrs0[5]
                }
        return gt_labels

    @staticmethod
    def get_rec_frms_raw(did:str, rid:str) -> Tuple[List]:
        '''
        获取当前未标注记录键值、图片、类别列表，通过键值可以更新gt_label中的内容
        '''
        rec_frms = []
        imgs = []
        folder = f'{DhlpWeb.base_folder}/{did}/{rid}/raw'
        index_fn = f'{DhlpWeb.base_folder}/{did}/{rid}/raw/index.txt'
        with open(index_fn, 'r', encoding='utf-8') as rfd:
            for row in rfd:
                row = row.strip()
                arrs0 = row.split(',')
                rec_frms.append({
                    'key': arrs0[0],
                    'image': f'{arrs0[1]}',
                    'folder': f'{DhlpWeb.base_folder}/{did}/{rid}/raw',
                    'cls': 'x'
                })
        return rec_frms

app = Flask(__name__)
app.config['UPLOAD_FOLDER'] = 'uploads'
app.config['RESULT_FOLDER'] = 'static/results'
# 创建目录
os.makedirs(app.config['UPLOAD_FOLDER'], exist_ok=True)
os.makedirs(app.config['RESULT_FOLDER'], exist_ok=True)

@app.route('/')
def index():
    DhlpWeb.gt_labels = DhlpWeb.get_gt_labels(DhlpWeb.ds_fn)
    return render_template('index.html')

@app.route('/get_curr_unlabel')
def get_curr_unlabel():
    print("进入 /get_curr_unlabel 路由")
    if not DhlpWeb.gt_labels:
        print("gt_labels 为空，开始加载")
        DhlpWeb.gt_labels = DhlpWeb.get_gt_labels(DhlpWeb.ds_fn)
    print("开始获取当前未标注记录")
    did, rid = DhlpWeb.get_curr_unlabel(DhlpWeb.labeling_prog_fn, DhlpWeb.labeling_progs)
    print(f"获取到未标注记录: did={did}, rid={rid}")
    rec_frms = DhlpWeb.get_rec_frms_raw(did, rid)
    for rec_frm in rec_frms:
        rec_frm['cls'] = DhlpWeb.gt_labels[rec_frm['key']]['cls']
    for rec_frm in rec_frms:
        print(f'### {rec_frm["key"]}: {rec_frm["image"]} , {rec_frm["cls"]};')
    resp = {
        'state': 1,
        'rec_frms': rec_frms
    }
    return jsonify(resp)

@app.route('/save_labels', methods=['POST'])
def save_labels():
    labels = request.get_json()
    for key, cls in labels.items():
        if key in DhlpWeb.gt_labels:
            DhlpWeb.gt_labels[key]['cls'] = cls

    # 更新标注进度
    did = DhlpWeb.gt_labels[list(labels.keys())[0]]['did']
    rid = DhlpWeb.gt_labels[list(labels.keys())[0]]['rid']
    key = f'{did}+{rid}'
    if key in DhlpWeb.labeling_progs:
        DhlpWeb.labeling_progs[key] = '1'
        with open(DhlpWeb.labeling_prog_fn, 'w', encoding='utf-8') as wfd:
            for k, v in DhlpWeb.labeling_progs.items():
                wfd.write(f'{k}:{v}\n')

    # 更新ds.txt文件
    with open(DhlpWeb.ds_fn, 'w', encoding='utf-8') as wfd:
        for k, v in DhlpWeb.gt_labels.items():
            wfd.write(f'{v["pt_fn"]},{v["did"]},{v["rid"]},{v["label"]},{v["pos"]},{v["cls"]}\n')

    resp = {
        'state': 1
    }
    return jsonify(resp)

# 假设你的图片存储在 static/images 目录下
@app.route('/images/<filename>')
def get_image(filename):
    folder = request.args.get('folder', type=str)
    print(f'######## fn: {filename}; folder: {folder}; ???????????????/')
    # 获取当前脚本所在目录
    script_dir = os.path.dirname(os.path.abspath(__file__))
    # 计算图片所在的绝对路径
    abs_folder = os.path.abspath(os.path.join(script_dir, folder))
    file_path = os.path.join(abs_folder, filename)
    if os.path.isfile(file_path):
        return send_from_directory(abs_folder, filename)
    else:
        return jsonify({"error": "File not found"}), 404

def t001():
    print(f'测试')
    # 找到第一个未标注设备+记录，开始进行标注
    did, rid = DhlpWeb.get_curr_unlabel(DhlpWeb.labeling_prog_fn)
    print(f'获取当前未标注记录：{did}, {rid}')

    print(f'^_^ The End! ^_^')

    

if __name__ == '__main__':
    try:
        app.run(host='0.0.0.0', port=8080, use_reloader=False, debug=True)
    except Exception as e:
        print(f"服务器启动失败: {e}")
    # t001()